#include "cmd_handler.h"

#include <windows.h>
#include <Tlhelp32.h>

#define CMD_PS		"ps"
#define CMD_LS		"ls"
#define CMD_CAT		"cat"
#define CMD_CP		"cp"
#define CMD_RM		"rm"
#define CMD_KILL	"kill"
#define CMD_FREE	"free"
#define CMD_HELP	"help"

// strip leading spaces
char* strip_before(const char* p)
{
	char* r = p;
	while(r[0] == ' ')
	{
		++r;
	}
	return r;
}

static void write_line(SOCKET sock, const char* line)
{
	send(sock, line, strlen(line), 0);
}

static void help_cmd_handler(const char* arg, SOCKET sock)
{
	const char* buf = 
		"ls      - list file\r\n"
		"cp      - copy file or directory\r\n"
		//"rm      - remove file or directory\r\n"
		"ps      - list process (use -t option to see more)\r\n"
		"kill    - kill a process\r\n"
		"cat     - read and output file\r\n"
		"free    - show memory usage\r\n"
		"command - run command\r\n";
	write_line(sock, buf);
	return;
}

static void ps_cmd_handler(const char* arg, SOCKET sock)
{
	int index = 0;
	char cw[1024];
	char tmp[1024];
	BOOL lRes;
	PROCESSENTRY32 procEntry;
	THREADENTRY32 thrdEntry;
	MODULEENTRY32 moduEntry;

	FILETIME ct, et, kt, ut;
	HANDLE lProcSnapshot = CreateToolhelp32Snapshot(TH32CS_SNAPALL,0);

	sprintf_s(tmp, sizeof(tmp), "pid\t\texe\t\tthread count\r\n", strlen("pid\t\texe\t\tthread count\r\n"));
	write_line(sock, tmp);
	if(lProcSnapshot != INVALID_HANDLE_VALUE)
	{
		memset(&procEntry, 0, sizeof(PROCESSENTRY32));
		procEntry.dwSize = sizeof(PROCESSENTRY32);
		lRes = Process32First(lProcSnapshot, &procEntry);
		while(lRes)
		{
			memset(cw, 0, sizeof(cw));
			WideCharToMultiByte(CP_ACP, 0, procEntry.szExeFile, wcslen(procEntry.szExeFile), cw, sizeof(cw)-1, NULL, NULL);
			sprintf_s(tmp, sizeof(tmp), "%d\t%s\t\t%d\r\n", procEntry.th32ProcessID, cw, procEntry.cntThreads);
			write_line(sock, tmp);
			lRes = Process32Next(lProcSnapshot, &procEntry);
		}

		if(0 == strncmp(strip_before(arg), "-t", strlen("-t")))
		{
			sprintf_s(tmp, sizeof(tmp), "\r\n**************** thread information ********** \r\ntid\t\towner pid\t\tkernel time\tuser time\r\n");
			write_line(sock, tmp);
			memset(&thrdEntry, 0, sizeof(THREADENTRY32));
			thrdEntry.dwSize = sizeof(THREADENTRY32);
			lRes = Thread32First(lProcSnapshot, &thrdEntry);
			while(lRes)
			{
				GetThreadTimes((HANDLE)(thrdEntry.th32ThreadID), &ct, &et, &kt, &ut);
				sprintf_s(tmp, sizeof(tmp), "%d\t%d\t\t{%d %d}\t{%d %d}\r\n", thrdEntry.th32ThreadID, thrdEntry.th32OwnerProcessID, 
					kt.dwHighDateTime, kt.dwLowDateTime, ut.dwHighDateTime, ut.dwLowDateTime);
				write_line(sock, tmp);
				lRes = Thread32Next(lProcSnapshot, &thrdEntry);

			}
		}

		//   memset(&moduEntry, 0, sizeof(MODULEENTRY32));
		//   moduEntry.dwSize = sizeof(MODULEENTRY32);
		//lRes = Module32First(lProcSnapshot, &moduEntry);
		//   while(lRes)
		//   {
		//	wprintf(TEXT("pid: %d  mid: %d  mod: %s path: %s\r\n"), moduEntry.th32ProcessID, moduEntry.th32ModuleID, moduEntry.szModule, moduEntry.szExePath);
		//	//wprintf(TEXT("%s\r\n"), thrdEntry. .szExeFile);
		//     //Use szExeFile member to get the module name of the current process.
		//     lRes = Module32Next(lProcSnapshot, &moduEntry);

		//   }
	}
	CloseToolhelp32Snapshot(lProcSnapshot);
}

static void kill_cmd_handler(const char* arg, SOCKET sock)
{
	int pid = atoi(arg);
	HANDLE hproc = OpenProcess(0, 0, pid);
	TerminateProcess(hproc, 0);
	CloseHandle(hproc);
}

static void free_cmd_handler(const char* arg, SOCKET sock)
{
	char buf[1024];
	MEMORYSTATUS mem_status;
	SYSTEM_INFO sys_info;

	memset(&mem_status, 0, sizeof(MEMORYSTATUS));
	mem_status.dwLength = sizeof(MEMORYSTATUS);
	GlobalMemoryStatus(&mem_status);
	memset(&sys_info, 0, sizeof(SYSTEM_INFO));
	GetSystemInfo(&sys_info);

	sprintf_s(buf, sizeof(buf), 
		"GlobalMemoryStatus:\r\n"
		"dwMemoryLoad: %d%%\r\n"
		"dwAvailPhys: %d\r\n"
		"dwTotalPhys: %d\r\n"
		"dwPageSize: %d\r\n"
		"dwAvailPageFile: %d\r\n"
		"dwTotalPageFile: %d\r\n",
		mem_status.dwMemoryLoad, mem_status.dwAvailPhys, mem_status.dwTotalPhys, 
		sys_info.dwPageSize, mem_status.dwAvailPageFile, mem_status.dwTotalPageFile
		);

	write_line(sock, buf);
}

static void ls_cmd_handler(const char* arg, SOCKET sock)
{
	TCHAR path[1024] = {0};
	char buf[1024];
	DWORD attr;
	WIN32_FIND_DATA find_data;
	HANDLE hfind;
	if(strlen(strip_before(arg)) < 1)
		arg = "\\";

	mbstowcs(path, strip_before(arg), sizeof(path));
	attr = GetFileAttributes(path);
	if(attr & FILE_ATTRIBUTE_DIRECTORY)
		wcscat_s(path, sizeof(path), TEXT("\\*"));

	hfind = FindFirstFile(path, &find_data);
	if(hfind != INVALID_HANDLE_VALUE)
	{
		do
		{
			memset(buf, 0, sizeof(buf));
			WideCharToMultiByte(CP_ACP, 0, find_data.cFileName, wcslen(find_data.cFileName), buf, sizeof(buf)-1, NULL, NULL);
			if(find_data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
				strcat_s(buf, sizeof(buf), "\\\r\n");
			else
				strcat_s(buf, sizeof(buf), "\r\n");
			write_line(sock, buf);
		} while(FindNextFile(hfind, &find_data));
	}
	//FindFirstFile();
}

static void cat_cmd_handler(const char* arg, SOCKET sock)
{
	FILE *f;
	int rc = 0, i, len = strlen(arg);
	char path[512];
	char buf[1024];

	arg = strip_before(arg);
	strncpy_s(path, sizeof(path), arg, len);
	for(i = 0; i < len; ++i)
	{
		if(path[i] == '\r' || path[i] == '\n')
		{
			path[i] = 0;
			break;
		}
	}
	rc = fopen_s(&f, path, "r");
	if(rc != 0)
	{
		sprintf_s(buf, sizeof(buf), "fopen %s failed, error is %d\r\n", path, GetLastError());
		printf(buf);
		write_line(sock, buf);
		return;
	}
	while(fread_s(buf, sizeof(buf), 1, 16, f))
		write_line(sock, buf);
}

static void run_cmd_handler(const char* arg, SOCKET sock)
{
	WCHAR path[1024];
	char buf[1024];
	mbstowcs(path, strip_before(arg), sizeof(path));
	if(!CreateProcess(path, NULL, NULL, NULL, NULL, 0, NULL, NULL, NULL, NULL))
	{
		sprintf_s(buf, sizeof(buf), "CreateProcess %s failed, error is %d\r\n", strip_before(arg), GetLastError());
		printf(buf);
		write_line(sock, buf);
		return;
	}
	sprintf_s(buf, sizeof(buf), "%s\r\n", strip_before(arg));
	printf(buf);
	write_line(sock, buf);
	return;
}

static void get_path(const char* arg, int* start, int* end)
{
	char* p = strip_before(arg);
	int has_d_quote = 0, len = strlen(p), i; 
	*start = 0, *end = len;

	if(p[0] == '"')
	{
		has_d_quote = 1;
		*start = 1;
	}
	
	for(i=*start; i < len; ++i)
	{
		if(has_d_quote)
		{
			if(p[i] == '"')
			{ // found matching double quote
				*end = i;
				break;
			}
		}
		else
		{
			if(p[i] == ' ')
			{
				*end = i;
				break;
			}
		}
	}
}

static void cp_cmd_handler(const char* arg, SOCKET sock)
{
	WCHAR src[1024], dst[1024];
	char* p = strip_before(arg), tmp[1024] = {0};
	int has_d_quote = 0, start = 0, end = 0, len = strlen(p), i;
	
	get_path(p, &start, &end);
	strncpy(tmp, p+start, end-start);
	tmp[end-start] = 0;
	mbstowcs(src, tmp, sizeof(src));
	
	p = strip_before(p+end+1);
	get_path(p, &start, &end);
	strncpy(tmp, p+start, end-start);
	tmp[end-start] = 0;	
	mbstowcs(dst, tmp, sizeof(dst));

	CopyFile(src, dst, 1);	
}

static void rm_cmd_handler(const char* arg, SOCKET sock)
{
	WCHAR path[1024];
	DWORD attr;
	mbstowcs(path, strip_before(arg), sizeof(path));
	attr = GetFileAttributes(path);
	if(attr & FILE_ATTRIBUTE_DIRECTORY)
	{
		if(RemoveDirectory(path) == 0)
		{
			printf("RemoveDirectory failed, error: %d", GetLastError());
		}
	}
	else
	{
		if(DeleteFile(path) == 0)
		{
			printf("DeleteFile failed, error: %d", GetLastError());
		}
	}
}

int exec_cmd(const char* cmd, SOCKET sock)
{
	const char* p = NULL;
	int id = 0;

	printf("exec_cmd : %s\r\n", cmd);

	if(strncmp(cmd, CMD_HELP, strlen(CMD_HELP)) == 0)
	{
		help_cmd_handler(cmd + strlen(CMD_HELP), sock);
		return 0;
	}

	if(strncmp(cmd, CMD_PS, strlen(CMD_PS)) == 0)
	{
		ps_cmd_handler(cmd + strlen(CMD_PS), sock);
		return 0;
	}

	if(strncmp(cmd, CMD_LS, strlen(CMD_LS)) == 0)
	{
		ls_cmd_handler(cmd + strlen(CMD_LS), sock);
		return 0;
	}

	if(strncmp(cmd, CMD_CAT, strlen(CMD_CAT)) == 0)
	{
		cat_cmd_handler(cmd + strlen(CMD_CAT), sock);
		return 0;
	}

	if(strncmp(cmd, CMD_KILL, strlen(CMD_KILL)) == 0)
	{
		kill_cmd_handler(cmd + strlen(CMD_KILL), sock);
		return 0;
	}

	if(strncmp(cmd, CMD_FREE, strlen(CMD_FREE)) == 0)
	{
		free_cmd_handler(cmd + strlen(CMD_FREE), sock);
		return 0;
	}

	if(strncmp(cmd, CMD_CP, strlen(CMD_CP)) == 0)
	{
		cp_cmd_handler(cmd + strlen(CMD_CP), sock);
		return 0;
	}

	if(strncmp(cmd, CMD_RM, strlen(CMD_RM)) == 0)
	{
		rm_cmd_handler(cmd + strlen(CMD_RM), sock);
		return 0;
	}

	// default to run application
	run_cmd_handler(cmd, sock);
}